package com.yf.tool.registrar;

import com.yf.constants.AiConstants;
import com.yf.model.ai.entity.AiConversationToolBinding;
import com.yf.model.ai.entity.AiTool;
import com.yf.model.ai.enums.AiToolTypeEnum;
import com.yf.service.manage.IAiConversationToolBindingService;
import com.yf.service.manage.IAiToolService;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.mcp.SyncMcpToolCallback;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * AI 能力 注册器
 *
 * @author : YiFei
 * @since : 2025/7/4 12:18
 */
@Component
@RequiredArgsConstructor
public class AiToolRegistrar {

    private final IAiToolService aiToolService;
    private final IAiConversationToolBindingService aiConversationToolBindingService;
    private final ToolCallbackProvider toolCallbackProvider;
    private final List<ToolCallback> toolCallbacks;

    /**
     * 同步 MCP 最新工具
     */
    @PostConstruct
    public void init() {
        // 同步 MCP 工具名
        syncMcpTools();
        // 同步 本地 Tools 方法
        syncToolCalling();
    }

    /**
     * 通用同步工具方法
     */
    private void syncTools(
            ToolCallback[] toolCallbacks,
            AiToolTypeEnum typeEnum,
            java.util.function.Consumer<AiTool> extraSetter
    ) {
        Set<String> toolNames = Arrays.stream(toolCallbacks)
                .map(tool -> tool.getToolDefinition().name())
                .collect(Collectors.toSet());

        List<AiTool> dbTools = aiToolService.lambdaQuery()
                .eq(AiTool::getComponentType, typeEnum)
                .list();
        Map<String, Integer> dbToolNameToId = dbTools.stream()
                .collect(Collectors.toMap(AiTool::getName, AiTool::getId));
        Set<String> dbToolNames = dbToolNameToId.keySet();

        Set<String> toAdd = new java.util.HashSet<>(toolNames);
        toAdd.removeAll(dbToolNames);

        Set<String> toDelete = new java.util.HashSet<>(dbToolNames);
        toDelete.removeAll(toolNames);

        // 批量新增
        if (!toAdd.isEmpty()) {
            List<AiTool> toolsToAdd = Arrays.stream(toolCallbacks)
                    .filter(tool -> toAdd.contains(tool.getToolDefinition().name()))
                    .map(tool -> {
                        AiTool aiTool = new AiTool();
                        aiTool.setName(tool.getToolDefinition().name());
                        aiTool.setDescription(tool.getToolDefinition().description());
                        aiTool.setComponentType(typeEnum);
                        aiTool.setIsActive(Boolean.TRUE);
                        if (extraSetter != null) {
                            extraSetter.accept(aiTool);
                        }
                        return aiTool;
                    })
                    .collect(Collectors.toList());
            aiToolService.saveBatch(toolsToAdd);
        }

        // 批量删除
        if (!toDelete.isEmpty()) {
            List<Integer> idsToDelete = toDelete.stream()
                    .map(dbToolNameToId::get)
                    .filter(java.util.Objects::nonNull)
                    .collect(Collectors.toList());
            aiToolService.removeBatchByIds(idsToDelete);
            aiConversationToolBindingService.lambdaUpdate()
                    .in(AiConversationToolBinding::getToolId, idsToDelete)
                    .remove();
        }
    }

    /**
     * 同步本地工具调用方法
     * 目前仅用于测试，实际使用中可以通过注入的方式获取工具回调
     */
    private void syncToolCalling() {
        ToolCallback[] toolCalling = toolCallbacks.toArray(new ToolCallback[0]);
        syncTools(toolCalling, AiToolTypeEnum.TOOL_CALLING, aiTool -> {
            aiTool.setProvider(AiConstants.AI_DEFAULT_PROVIDER);
            aiTool.setProviderIcon(AiConstants.AI_DEFAULT_PROVIDER_ICON);
        });
    }

    /**
     * 同步 Mcp 方法
     */
    private void syncMcpTools() {
        ToolCallback[] mcpTools = toolCallbackProvider.getToolCallbacks();
        syncTools(mcpTools, AiToolTypeEnum.MCP, null);
    }

    /**
     * 获取 AI 工具回调列表
     *
     * @param aiToolNames AI 工具名称集合
     * @return 工具回调列表
     */
    public List<ToolCallback> getToolCallbacks(Set<String> aiToolNames) {
        ToolCallback[] mcpTools = toolCallbackProvider.getToolCallbacks();

        List<ToolCallback> toolCallbackList = new ArrayList<>(mcpTools.length + toolCallbacks.size());

        toolCallbackList.addAll(toolCallbacks);
        toolCallbackList.addAll(List.of(mcpTools));

        return toolCallbackList.stream()
                .filter(toolCallback -> aiToolNames.contains(toolCallback.getToolDefinition().name()))
                .toList();
    }

}
